{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## KV Cache\n",
    "\n",
    "只用在推理过程中。\n",
    "\n",
    "“我是”-->“周”\n",
    "\n",
    "“我是周”--->“瑜”\n",
    "\n",
    "“我是周瑜”--->“...”"
   ],
   "id": "fb712db5ee1d4280"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-27T09:31:17.029833Z",
     "start_time": "2025-07-27T09:31:17.020757Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim: int, attn_dim: int, output_dim: int, num_heads: int):\n",
    "        super().__init__()\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.attn_dim = attn_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = attn_dim // num_heads # //表示向下取整，attn_dim是head_dim的整数倍\n",
    "\n",
    "        # QKV投影层：从输入维度映射到内部维度\n",
    "        # projection\n",
    "        self.q_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "        self.k_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "        self.v_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "\n",
    "        # 输出投影层：从内部维度映射到输出维度\n",
    "        self.out_proj = nn.Linear(self.attn_dim, self.output_dim)\n",
    "\n",
    "    def forward(self, x, kv_cache=None):\n",
    "        \"\"\"\n",
    "        输入: [batch_size, seq_len, embed_dim]\n",
    "        返回: [batch_size, seq_len, output_dim]\n",
    "        \"\"\"\n",
    "        batch_size, seq_len, embed_dim = x.shape\n",
    "\n",
    "        # 投影到QKV空间\n",
    "        q = self.q_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "        k = self.k_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "        v = self.v_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        # 分割多头 [batch_size, num_heads, seq_len, head_dim]\n",
    "        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        if kv_cache is not None:\n",
    "            # 获取缓存的kv\n",
    "            k_cache, v_cache = kv_cache\n",
    "            # 拼接缓存的kv和当前的kv\n",
    "            k = torch.cat([k_cache, k], dim=2)  # 256, 128 --> cache ---> 128, 256  MLA\n",
    "            v = torch.cat([v_cache, v], dim=2)\n",
    "            # 更新缓存的kv\n",
    "            kv_cache = (k, v)\n",
    "        else:\n",
    "            kv_cache = (k, v)\n",
    "\n",
    "        # 计算注意力得分\n",
    "        # q   [batch_size, num_heads, seq_len, head_dim]\n",
    "        # k.T [batch_size, num_heads, head_dim, seq_len]\n",
    "        # q @ k.T 形状: [batch_size, num_heads, seq_len, seq_len]\n",
    "        attn_scores = torch.matmul(q, k.transpose(-2, -1))\n",
    "\n",
    "        # 缩放因子：防止乘积过大\n",
    "        d_k = k.size(-1)\n",
    "        attn_scores = attn_scores / torch.sqrt(torch.tensor(d_k))\n",
    "\n",
    "        # 计算注意力权重\n",
    "        attn_weights = torch.softmax(attn_scores, dim=-1)\n",
    "\n",
    "        # 计算注意力输出\n",
    "        # attn_weights [batch_size, num_heads, seq_len, seq_len]\n",
    "        # v            [batch_size, num_heads, seq_len, head_dim]\n",
    "        # [batch_size, num_heads, seq_len, head_dim]\n",
    "        attn_out = torch.matmul(attn_weights, v)\n",
    "\n",
    "        # 合并多头 [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        attn_out = attn_out.transpose(1, 2).reshape(batch_size, seq_len, self.attn_dim)\n",
    "\n",
    "        # 投影到输出空间\n",
    "        return self.out_proj(attn_out), kv_cache"
   ],
   "id": "403c5a8fcae45aae",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-27T09:35:42.108167Z",
     "start_time": "2025-07-27T09:31:18.974805Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 一个seq\n",
    "x = torch.rand(1, 4, 2)\n",
    "\n",
    "attn = MultiHeadAttention(embed_dim=2, attn_dim=6, output_dim=8, num_heads=2)\n",
    "out, kv_cache = attn(x)\n",
    "\n",
    "print(x.shape)\n",
    "print(out.shape)\n",
    "\n",
    "x = torch.rand(1, 5, 2)\n",
    "out, kv_cache = attn(x[:, -1:, :], kv_cache)\n",
    "print(x.shape)\n",
    "print(out.shape)"
   ],
   "id": "6c229fa626c5ff1c",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 2])\n",
      "torch.Size([1, 4, 8])\n",
      "torch.Size([1, 5, 2])\n",
      "torch.Size([1, 1, 8])\n"
     ]
    }
   ],
   "execution_count": 4
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
